package org.elasticsearch.index.analysis;

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 * <p/>
 * http://www.apache.org/licenses/LICENSE-2.0
 * <p/>
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import org.apache.lucene.analysis.TokenFilter;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.analysis.tokenattributes.CharTermAttribute;
import org.apache.lucene.analysis.tokenattributes.PositionIncrementAttribute;
import org.elasticsearch.analysis.PinyinConfig;
import org.elasticsearch.commons.Pinyin;


import java.io.IOException;
import java.util.*;
import java.util.stream.Collectors;

public class MultiPinyinTokenFilter extends TokenFilter {

    private final CharTermAttribute termAtt = addAttribute(CharTermAttribute.class);
    private boolean done = true;
    private boolean processedCandidate = false;
    private boolean processedFullPinyinLetter = false;
    private boolean processedFirstLetter = false;
    private boolean processedOriginal = false;
    private boolean processedSortCandidate = false;
    protected int position = 0;
    protected int lastOffset = 0;
    private PinyinConfig config;
    List<TermItem> candidate;
    private HashSet<String> termsFilter;

    protected int candidateOffset = 0;
    StringBuilder firstLetters;

    LinkedHashMap<String, List<String>> firstLetterMaps;
    StringBuilder fullPinyinLetters;

    LinkedHashMap<String, List<String>> fullPinyinLetterMaps;
    String source;
    private int lastIncrementPosition = 0;

    private PositionIncrementAttribute positionAttr = addAttribute(
            PositionIncrementAttribute.class);

    public MultiPinyinTokenFilter(TokenStream in, PinyinConfig config) {
        super(in);
        this.config = config;
        //validate config
        if (!(config.keepFirstLetter || config.keepSeparateFirstLetter || config.keepFullPinyin
                || config.keepJoinedFullPinyin || config.keepSeparateChinese)) {
            throw new ConfigErrorException(
                    "pinyin config error, can't disable separate_first_letter, first_letter and full_pinyin at the same time.");
        }
        candidate = new ArrayList<>();
        firstLetters = new StringBuilder();
        termsFilter = new HashSet<>();
        fullPinyinLetters = new StringBuilder();
        firstLetterMaps = new LinkedHashMap<>();
        fullPinyinLetterMaps = new LinkedHashMap<>();
    }

    //TODO refactor, merge code
    @Override
    public final boolean incrementToken() throws IOException {

        if (!done) {
            if (readTerm()) {
                return true;
            }
        }

        if (done) {
            resetVariable();
            if (!input.incrementToken()) {
                return false;
            }
            done = false;
        }
        readTerm();
        return true;
    }

    private boolean readTerm() {
        if (!processedCandidate) {
            processedCandidate = true;
            lastOffset = termAtt.length();
            source = termAtt.toString();
            if (config.trimWhitespace) {
                source = source.trim();
            }

            Map<String, List<String>> pinyinListMaps = Pinyin.multiPinyin(source);
            List<String> chineseList = ChineseUtil.segmentChinese(source);
            if (pinyinListMaps.size() == 0 || chineseList.isEmpty()) {
                return false;
            }

            StringBuilder buff = new StringBuilder();
            int buffStartPosition = 0;
            int buffSize = 0;
            position = 0;
            for (int i = 0; i < source.length(); i++) {
                char c = source.charAt(i);

                //keep original alphabet
                // 当查询的为拼音时
                if (c < 128) {
                    if (buff.length() <= 0) {
                        buffStartPosition = i;
                    }
                    if ((c > 96 && c < 123) || (c > 64 && c < 91) || (c > 47 && c < 58)) {
                        if (config.keepNoneChinese) {
                            if (config.keepNoneChinese) {
                                if (config.keepNoneChineseTogether) {
                                    buff.append(c);
                                    buffSize++;
                                } else {
                                    addCandidate(new TermItem(String.valueOf(c), i, i + 1,
                                            buffStartPosition));
                                }
                            }
                        }
                        if (config.keepNoneChineseInFirstLetter) {
                            firstLetters.append(c);
                        }
                        if (config.keepNoneChineseInJoinedFullPinyin) {
                            fullPinyinLetters.append(c);
                        }
                    }
                } else {
                    //clean previous temp
                    if (buff.length() > 0) {
                        buffSize = parseBuff(buff, buffSize, buffStartPosition);
                    }
                    String chinese = chineseList.get(i);
                    List<String> pinyins = pinyinListMaps.getOrDefault(chinese, null);
                    if (pinyins != null && !pinyins.isEmpty()) {
                        position++;
                        List<String> firstLetterLists = pinyins.stream().map(s -> String.valueOf(s.charAt(0))).distinct().collect(Collectors.toList());
                        firstLetterMaps.put(chinese, firstLetterLists);
                        if (config.keepJoinedFullPinyin) {
                            fullPinyinLetterMaps.put(chinese, pinyins);
                        }
                        if (config.keepSeparateChinese) {
                            addCandidate(new TermItem(chinese, i, i + 1, position));
                        }
                    }
                }

                lastOffset = i;

            }

            //clean previous temp
            if (buff.length() > 0) {
                buffSize = parseBuff(buff, buffSize, buffStartPosition);
            }
        }

        if (config.keepOriginal && !processedOriginal) {
            processedOriginal = true;
            addCandidate(new TermItem(source, 0, source.length(), 1));
        }

        if (config.keepJoinedFullPinyin && !processedFullPinyinLetter && fullPinyinLetterMaps != null && !fullPinyinLetterMaps.isEmpty()) {
            processedFullPinyinLetter = true;

            List<String> combinations = combineLists(fullPinyinLetterMaps, new ArrayList<>(fullPinyinLetterMaps.keySet()));
            for (String fullPinyinLetter : combinations) {
                addCandidate(new TermItem(fullPinyinLetter, 0, source.length(), 1));
            }
            fullPinyinLetterMaps.clear();
        }

        if (config.keepFirstLetter && firstLetterMaps != null && firstLetterMaps.size() > 0 && !processedFirstLetter) {
            processedFirstLetter = true;
            List<String> combinations = combineLists(firstLetterMaps, new ArrayList<>(firstLetterMaps.keySet()));

            for (String firstLetters : combinations) {
                String fl;
                if (firstLetters.length() > config.LimitFirstLetterLength
                        && config.LimitFirstLetterLength > 0) {
                    fl = firstLetters.substring(0, config.LimitFirstLetterLength);
                } else {
                    fl = firstLetters;
                }
                if (config.lowercase) {
                    fl = fl.toLowerCase();
                }
                if (!(config.keepSeparateFirstLetter && fl.length() <= 1)) {
                    addCandidate(new TermItem(fl, 0, fl.length(), 1));
                }
            }
            firstLetterMaps.clear();
        }

        if (!processedSortCandidate) {
            processedSortCandidate = true;
            Collections.sort(candidate);
        }

        if (candidateOffset < candidate.size()) {
            TermItem item = candidate.get(candidateOffset);
            candidateOffset++;
            setTerm(item.term, item.startOffset, item.endOffset, item.position);
            return true;
        }

        done = true;
        return false;
    }

    public static List<String> combineLists(Map<String, List<String>> map,
                                            List<String> keysToCombine) {
        List<String> result = new ArrayList<>();

        // 验证所有key是否存在于map中
        for (String key : keysToCombine) {
            if (!map.containsKey(key)) {
                throw new IllegalArgumentException("Key " + key + " not found in the map.");
            }
        }
        // 使用递归或迭代生成所有组合
        backtrack(map, keysToCombine, new ArrayList<>(), 0, result);

        return result;
    }

    private static void backtrack(Map<String, List<String>> map, List<String> keys,
                                  List<String> currentCombination, int index, List<String> result) {
        if (index == keys.size()) {
            // 所有key的列表都已经被遍历过，将当前组合添加到结果中
            result.add(String.join("", currentCombination));
            return;
        }
        String key = keys.get(index);
        List<String> list = map.get(key);
        // 遍历当前key对应的列表
        for (String item : list) {
            currentCombination.add(item); // 将当前元素添加到当前组合中
            backtrack(map, keys, currentCombination, index + 1, result); // 递归处理下一个key
            currentCombination.remove(currentCombination.size() - 1); // 回溯，移除刚刚添加的元素以尝试下一个元素
        }
    }

    void addCandidate(TermItem item) {

        String term = item.term;
        if (config.lowercase) {
            term = term.toLowerCase();
        }

        if (config.trimWhitespace) {
            term = term.trim();
        }
        item.term = term;

        if (term.length() == 0) {
            return;
        }

        //remove same term with same position
        String fr = term + item.position;

        //remove same term, regardless position
        if (config.removeDuplicateTerm) {
            fr = term;
        }

        if (termsFilter.contains(fr)) {
            return;
        }
        termsFilter.add(fr);

        candidate.add(item);
    }


    void setTerm(String term, int startOffset, int endOffset, int position) {
        if (config.lowercase) {
            term = term.toLowerCase();
        }

        if (config.trimWhitespace) {
            term = term.trim();
        }

        //ignore empty term
        if (term.length() == 0) {
            return;
        }

        termAtt.setEmpty();
        termAtt.append(term);
        if (startOffset < 0) {
            startOffset = 0;
        }
    }


    private int parseBuff(StringBuilder buff, int buffSize, int buffPosition) {
        if (config.keepNoneChinese) {
            if (config.noneChinesePinyinTokenize) {
                List<String> result = PinyinAlphabetTokenizer.walk(buff.toString());
                int start = (lastOffset - buffSize + 1);
                for (int i = 0; i < result.size(); i++) {
                    int end;
                    String t = result.get(i);
                    if (config.fixedPinyinOffset) {
                        end = start + 1;
                    } else {
                        end = start + t.length();
                    }
                    addCandidate(new TermItem(result.get(i), start, end, ++position));
                    start = end;
                }
            } else if (config.keepFirstLetter || config.keepSeparateFirstLetter
                    || config.keepFullPinyin || !config.keepNoneChineseInJoinedFullPinyin) {
                addCandidate(
                        new TermItem(buff.toString(), lastOffset - buffSize, lastOffset, ++position));
            }
        }

        buff.setLength(0);
        buffSize = 0;
        return buffSize;
    }

    @Override
    public final void end() throws IOException {
        super.end();
    }

    void resetVariable() {
        position = 0;
        lastOffset = 0;
        candidate.clear();
        this.processedCandidate = false;
        this.processedFirstLetter = false;
        this.processedFullPinyinLetter = false;
        this.processedOriginal = false;
        firstLetters.setLength(0);
        fullPinyinLetters.setLength(0);
        source = null;
        candidateOffset = 0;
        termsFilter.clear();
        lastIncrementPosition = 0;
    }

    @Override
    public void reset() throws IOException {
        super.reset();
        this.done = true;
        resetVariable();
    }


}
